import os
import pandas as pd
from tqdm import tqdm
from collections import defaultdict
import pyarrow # 核心诊断：强制导入pyarrow

def generate_buy_signals_with_pandas(all_stock_codes, market_df, cache_dir, start_date_str, end_date_str, momentum_threshold, max_pullback_pct, max_price, required_start_date_obj, ma_divergence_20_max, ma_divergence_60_max, rsi_overbought_threshold, precise_pullback_max_dist_pct, market_rsi_max):
    """使用Pandas进行快速的向量化计算，生成回测期间内所有的买入信号。"""
    print("开始加载所有缓存数据到内存...")
    all_dfs = []
    for code in tqdm(all_stock_codes, desc="加载数据"):
        # 修复：与data_provider中逻辑统一，智能处理股票代码前缀
        if 'sh.' in code or 'sz.' in code:
            symbol = code
        else:
            symbol = f"sh.{code}" if code.startswith('6') else f"sz.{code}"

        fpath = os.path.join(cache_dir, f"{symbol}.feather")
        
        if os.path.exists(fpath):
            stock_df = pd.read_feather(fpath)
            # 此处原有的次新股过滤器已被移除，因为过滤操作已前置
            stock_df['symbol'] = symbol
            all_dfs.append(stock_df)

    if not all_dfs:
        print("警告：未能加载任何股票数据进行海选。")
        return {}, []

    df = pd.concat(all_dfs, ignore_index=True)
    df['date'] = pd.to_datetime(df['date'])
    
    # 注意：这里的日期过滤只应在所有指标计算完毕后进行，但为了海选效率和内存管理，
    # 我们先不过滤，让指标在包含预热期的完整数据上计算。
    # 最终的信号筛选会自然落在回测区间内。

    print(f"数据加载完成，共计 {df['symbol'].nunique()} 只股票，{len(df)} 条记录。开始计算指标...")
    market_df['market_ma60'] = market_df['close'].rolling(60).mean()
    market_df['market_ok'] = market_df['close'] > market_df['market_ma60']
    market_df['market_is_up'] = market_df['close'] > market_df['close'].shift(1)
    
    # 新增：为大盘计算RSI
    market_delta = market_df['close'].diff()
    market_gain = market_delta.where(market_delta > 0, 0).ewm(span=14, adjust=False).mean()
    market_loss = -market_delta.where(market_delta < 0, 0).ewm(span=14, adjust=False).mean()
    market_rs = market_gain / market_loss
    market_df['market_rsi'] = 100 - (100 / (1 + market_rs))

    df = pd.merge(df, market_df[['date', 'market_ok', 'market_is_up', 'market_rsi']], on='date', how='left')
    df['market_ok'] = df['market_ok'].fillna(False)
    df['market_is_up'] = df['market_is_up'].fillna(False)
    df['market_rsi'] = df['market_rsi'].fillna(50) # 如果没有值，给一个中性值

    df.set_index(['symbol', 'date'], inplace=True)
    df.sort_index(inplace=True)
    
    # --- RSI计算 ---
    delta = df.groupby(level='symbol')['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.groupby(level='symbol').transform(lambda x: x.ewm(span=14, adjust=False).mean())
    avg_loss = loss.groupby(level='symbol').transform(lambda x: x.ewm(span=14, adjust=False).mean())
    rs = avg_gain / avg_loss
    df['rsi'] = 100 - (100 / (1 + rs))
    # --- RSI计算结束 ---

    for p in [5, 10, 20, 60, 250]: # 增加250日均线
        df[f'ma{p}'] = df.groupby(level='symbol')['close'].rolling(window=p).mean().reset_index(level=0, drop=True)
    df['high_in_20d'] = df.groupby(level='symbol')['high'].rolling(window=20).max().reset_index(level=0, drop=True)
    df['price_change_ratio'] = df.groupby(level='symbol')['close'].pct_change().fillna(0) + 1
    df['is_limit_up'] = (df['price_change_ratio'] >= 1.099)
    df['is_limit_down'] = (df['price_change_ratio'] <= 0.901)
    
    # 使用 .values 绕过 Pandas 在 groupby().rolling() 中复杂的索引对齐问题。
    # 这依赖于 DataFrame 的索引是预先排序的 (df.sort_index(inplace=True))。
    # .values 会返回一个 numpy 数组，其顺序与排过序的 DataFrame 一致。
    # Pandas 支持 Series 和同样长度的 numpy 数组之间的直接运算。
    sum_limit_up_2d = df.groupby(level='symbol')['is_limit_up'].rolling(window=2).sum().values
    sum_limit_down_2d = df.groupby(level='symbol')['is_limit_down'].rolling(window=2).sum().values
    df['no_extreme_volatility'] = (sum_limit_up_2d == 0) & \
                                (sum_limit_down_2d == 0) & \
                                (df['price_change_ratio'] > 0.91)

    df['has_limit_up_in_20d'] = (df.groupby(level='symbol')['is_limit_up'].rolling(window=20).sum().values >= 1)
    df['amount'] = df['volume'] * df['close']

    # 分别计算各项滚动指标的 .values
    vol_roll_8d_min = df.groupby(level='symbol')['volume'].rolling(window=8).min().values
    vol_shift1_roll_7d_min = df.groupby(level='symbol')['volume'].shift(1).rolling(window=7).min().values
    amount_roll_10d_max = df.groupby(level='symbol')['amount'].rolling(window=10).max().values
    
    # 直接在numpy数组层面进行比较，避免索引问题
    low_vol_1 = df['volume'].values <= vol_roll_8d_min
    low_vol_2 = df['volume'].values <= vol_shift1_roll_7d_min
    low_amount_cond = df['amount'].values < amount_roll_10d_max * 0.25

    # 合并条件
    df['volume_cond'] = low_vol_1 | low_vol_2 | low_amount_cond

    df['high_in_20d_prev1'] = df.groupby(level='symbol')['high_in_20d'].shift(1)
    df['close_prev5'] = df.groupby(level='symbol')['close'].shift(5)
    df['close_prev10'] = df.groupby(level='symbol')['close'].shift(10)
    df['close_prev60'] = df.groupby(level='symbol')['close'].shift(60)

    print("指标计算完成，开始筛选信号...")

    # --- 核心买入条件 Section ---

    # 新增：年线拐头向上条件
    df['ma250_prev1'] = df.groupby(level='symbol')['ma250'].shift(1)
    condition_ma250_upward = df['ma250'] > df['ma250_prev1']

    # 1. 动量条件 (Momentum Condition)
    # 计算5日回报倍数，例如，1.05代表5天上涨了5%
    momentum_5d = df['close'] / df['close_prev5']
    condition_momentum = momentum_5d > momentum_threshold

    # 2. 回撤条件 (Pullback Condition)
    # 计算从20日高点回撤的百分比
    pullback_from_20d_high = (df['close'] - df['high_in_20d_prev1']) / df['high_in_20d_prev1'] * 100
    # 确保回撤幅度在我们允许的范围内 (例如，>-12，意味着回撤不超过12%)
    condition_pullback = pullback_from_20d_high > max_pullback_pct

    # 3. 中期动量条件 (Mid-term Momentum)
    condition_mid_term_momentum = df['close'] > df['close_prev10']

    # 4. 均线多头排列
    condition_ma_align = (df['ma5'] > df['ma10']) & \
                         (df['ma10'] > df['ma20']) & \
                         (df['ma20'] > df['ma60'])
    
    # 5. 长期动量 (收盘价 > 60日前收盘价)
    condition_long_term_momentum = df['close'] > df['close_prev60']

    # 6. 短期均线位置 (收盘价 > MA10 且 > MA20)
    condition_short_term_pos = (df['close'] >= df['ma10']) & (df['close'] > df['ma20'])

    # 7. K线形态 (阴线)
    condition_is_yin_line = df['close'] < df['open']

    # 8. 日内回撤形态 (无长上影线)
    condition_no_long_upper_shadow = (df['high'] - df['close']) / df['high'] <= 0.03

    # 9. 均线乖离率条件 (防过热)
    ma10_div_ma20 = (df['ma10'] / df['ma20']) - 1
    ma10_div_ma60 = (df['ma10'] / df['ma60']) - 1
    condition_ma_divergence = (ma10_div_ma20 <= ma_divergence_20_max) & (ma10_div_ma60 <= ma_divergence_60_max)

    # 10. RSI防过热条件
    condition_not_overbought = df['rsi'] < rsi_overbought_threshold

    # --- 激活"僵尸"指标 ---
    # 11. 大盘择时条件 (大盘指数在60日线之上)
    condition_market_ok = df['market_ok']
    # 12. 市场顺风条件 (大盘当日上涨)
    # condition_market_tailwind = df['market_is_up']
    # 13. 波动率条件 (近期无连续涨跌停)
    condition_volatility = df['no_extreme_volatility']
    # 14. 强势股基因 (20日内曾有涨停)
    condition_had_limit_up = df['has_limit_up_in_20d']
    # 15. 成交量条件 (缩量)
    condition_volume = df['volume_cond']
    # 16. 价格上限过滤
    condition_max_price = df['close'] < max_price

    # 17. 新增：价格与5日线关系（收盘价必须低于或等于5日线）
    condition_close_under_ma5 = df['close'] >= df['ma5']

    # 18. 新增：趋势强度约束 (均线分离度)
    # 要求10日均线的值至少要比60日均线高出8%
    condition_trend_strength = (df['ma10'] / df['ma60'] - 1) > 0.08

    # 19. 新增：精准回调约束 (紧贴均线)
    # 要求收盘价与10日均线的差距在2%以内
    condition_precise_pullback = abs(df['close'] - df['ma10']) / df['ma10'] < precise_pullback_max_dist_pct

    # 20. 新增：大盘RSI过滤器
    condition_market_rsi_ok = df['market_rsi'] < market_rsi_max

    # 合并所有条件 (重构为更安全的格式)
    buy_conditions = (
        #condition_ma250_upward # 新增年线条件
        # condition_momentum
         condition_pullback
        & condition_mid_term_momentum
        & condition_ma_align
        & condition_long_term_momentum
        & condition_short_term_pos
        & condition_is_yin_line
        & condition_no_long_upper_shadow
        & condition_ma_divergence
        & condition_not_overbought
        & condition_market_ok
        # & condition_market_tailwind
        & condition_volatility
        & condition_had_limit_up
        & condition_volume
        & condition_max_price
        #& condition_close_under_ma5
        & condition_trend_strength
        #& condition_precise_pullback
        #& condition_market_rsi_ok
    )

    # 获取满足所有条件的日期
    buy_signals = df[buy_conditions]
    
    # 关键一步：只返回在回测时间窗口内的信号
    # 修正: 对 reset_index() 后的新DataFrame进行筛选，而不是原始的buy_signals
    buy_signals_with_cols = buy_signals.reset_index()
    signals_df = buy_signals_with_cols[
        (buy_signals_with_cols['date'] >= pd.to_datetime(start_date_str)) &
        (buy_signals_with_cols['date'] <= pd.to_datetime(end_date_str))
    ][['date', 'symbol']]

    if signals_df.empty:
        print("海选未发现任何满足条件的买入信号。")
        return {}, [] # 返回空的信号字典和股票列表

    # 将信号DataFrame转换为一个更高效的查找结构：{date: [symbol1, symbol2, ...]}
    buy_signals = defaultdict(list)
    for _, row in signals_df.iterrows():
        # 将Pandas的Timestamp转换为Python的datetime.date对象，以便在Backtrader中进行匹配
        signal_date = row['date'].date()
        buy_signals[signal_date].append(row['symbol'])

    hot_stocks = signals_df['symbol'].unique().tolist()
    print(f"信号生成完成，共发现 {len(hot_stocks)} 只股票在回测期间内产生了 {len(signals_df)} 个买入信号。")
    
    # 提取hot_stocks的原始数据，用于直接注入回测引擎，避免二次读取
    # df的索引是 ['symbol', 'date']
    hot_stocks_data = df.loc[hot_stocks]
    # 使用字典推导式更高效地创建数据字典
    hot_stocks_dfs = {symbol: group.reset_index().drop(columns='symbol') 
                      for symbol, group in hot_stocks_data.groupby(level='symbol')}

    return buy_signals, hot_stocks, hot_stocks_dfs 